import claripy

from ..shellcode import Shellcode


class X86SetRegister(Shellcode):

    os = ["cgc", "unix"]
    arches = ["X86"]
    name = "setregister"
    codes = {
            'eax': [b"\xb8", b"\xbb", b"\xff\xe3"],
            'ebx': [b"\xbb", b"\xb8", b"\xff\xee"],
            'ecx': [b"\xb9", b"\xbb", b"\xff\xe3"],
            'edx': [b"\xba", b"\xbb", b"\xff\xe3"],
            'edi': [b"\xbf", b"\xbb", b"\xff\xe3"],
            'esi': [b"\xbe", b"\xbb", b"\xff\xe3"],
            'esp': [b"\xbc", b"\xbb", b"\xff\xe3"],
            'ebp': [b"\xbd", b"\xbb", b"\xff\xe3"],
            }

    def __init__(self, register, value, pc):
        """
        :param register:    The register to set
        :param value:       The value to set register to
        :param pc:          Value of PC to crash at.
        """

        self.register = register
        self.value = value
        self.pc = pc

    def raw(self, arch=None):

        register = self.register
        value = self.value
        ip = self.pc

        if isinstance(value, int):
            value = claripy.BVV(value, 32)
        if isinstance(ip, int):
            ip = claripy.BVV(ip, 32)

        try:
            code_row = [claripy.BVV(x) for x in self.codes[register]]
        except KeyError:
            raise ValueError("register '%s' does not exist" % register)

        return claripy.Concat(code_row[0], value.reversed, code_row[1], ip.reversed,  code_row[2])
